'''

首先，我们导入了matplotlib库并添加了绘图的代码。在绘图代码块中，我们先创建了一个名为fig的figure对象和名为ax的axes对象，并设置了图表的大小。然后，使用ax.plot()
方法分别绘制了股价曲线、唐奇安通道的上下轨和20日均线，调用ax.legend()
方法添加图例，使用ax.set_title()
方法和ax.set_xlabel()
方法和ax.set_ylabel()
方法分别设置了图表的标题和轴标签。

最后，调用plt.show()
方法显示图表。
'''

import yfinance as yf
import pandas as pd
import matplotlib.pyplot as plt


def donchian_channel(ticker, period):
    df = yf.download(ticker, period=period)
    dc = pd.DataFrame(columns=['Date', 'Upper', 'Lower', '20MA'])
    for i in range(len(df) - (20 - 1)):
        high = max(df['High'][i:i + 20])
        low = min(df['Low'][i:i + 20])
        ma = df['Close'][i:i + 20].mean()
        dc = pd.concat(
            [dc, pd.DataFrame([[df.index[i + 20 - 1], high, low, ma]], columns=['Date', 'Upper', 'Lower', '20MA'])],
            axis=0, ignore_index=True)
    dc['Signal'] = ''
    for i in range(len(dc)):
        if dc['20MA'][i] > dc['Upper'][i]:
            dc['Signal'][i] = 'Sell'
        elif dc['20MA'][i] < dc['Lower'][i]:
            dc['Signal'][i] = 'Buy'
        else:
            dc['Signal'][i] = 'Hold'

    dc['Cross20MA'] = ''
    for i in range(len(dc)):
        if i == 0:
            continue
        if dc['20MA'][i] > dc['20MA'][i - 1]:
            dc['Cross20MA'][i] = 'Above'
        elif dc['20MA'][i] < dc['20MA'][i - 1]:
            dc['Cross20MA'][i] = 'Below'
        else:
            dc['Cross20MA'][i] = ''

    head = dc[['Date', '20MA', 'Upper', 'Lower', 'Signal', 'Cross20MA']]
    head.columns = ['日期', '20日均线', '唐奇安通道上轨', '唐奇安通道下轨', '信号', '20日均线突破情况']
    print(head)

    plt.style.use('ggplot')
    fig, ax = plt.subplots(figsize=(12, 6))
    ax.plot(df['Close'], label='Close')
    ax.plot(dc['Upper'], label='Upper', linewidth=1)
    ax.plot(dc['Lower'], label='Lower', linewidth=1)
    ax.plot(dc['20MA'], label='20MA', linestyle='--', linewidth=1)
    ax.legend(loc='upper left')
    ax.set_title(ticker)
    ax.set_xlabel('Date')
    ax.set_ylabel('Price')
    plt.show()

stock='601919.SS'
donchian_channel(stock, '2y')
